#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
    Module Documentation
    here
"""

# Created by  : Zhang Chengdong
# Create Date : 2025/2/27 13:58
# Version = v0.1.0

__author__ = "Zhang Chengdong"
__copyright__ = "Copyright 2024. Large scale model"
__credits__ = ['Zhang Chengdong']

__liscence__ = "MIT"
__version__ = "1.0.1"
__maintainer__ = "Zhang Chengdong"
__status__ = "Production"

import os
import requests
import yaml
import logging
from datetime import datetime
from functools import wraps
from typing import Optional, List, Dict, Tuple

from kubernetes import client, config

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, FileResponse

from model_config import DatabaseInfo
from service.data_preprocessing.data_processing import DataProcessing
from service.catboost_model.training_data import train_main as training_main

# ============================ 日志修饰器 ==========================
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def logging_decorator(func):
    """
    日志修饰器logger
    """

    @wraps(func)
    async def wrapper(*args, **kwargs):
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        func_name = func.__name__
        logging.info(f"[{current_time}] Function: {func_name} ")
        json_content = await func(*args, **kwargs)
        logging.info(f"[{current_time}] Function: {func_name} - Response: {json_content}")
        return json_content

    return wrapper


# ========================================== 服务初始化 ==================================
app = FastAPI()
DATABASEINFO: Optional[DatabaseInfo]


# ========================================= 获取所有的租户ID =============================
def get_all_tenant_all(url: str, header: dict = None) -> List[Tuple]:
    """
    :param url: 云平台环境，获取所有租户信息
    :param header: 请求头信息，静态
    :return: 返回租户id和租户名称
    """
    if header is None:
        response = requests.get(url)
    else:
        response = requests.get(url, headers=header)

    if response.status_code != 200:
        logging.info("获取所有租户ID失败！")
        return []
    data = response.json()
    tenant_id_name_list = []
    for item in data['data']:
        if item['customerName'] is None:
            continue
        if item['superiorCustomerName'] is None:
            continue
        if "非标产品区域" in item['superiorCustomerName']:
            continue
        if "pro" in item['customerEnvCode']:
            continue
        tenant_id_name_list.append((item['tenantId'], item['customerName']))
    return tenant_id_name_list


def get_config_of_k8s(choice_config: dict) -> Dict:
    """
    获取pod中的相关信息
    :param choice_config:
    :return:
    """
    if 'KUBERNETES_SERVICE_HOST' in os.environ:
        namespace = "taihang"
        # TODO 服务上线需要调整服务的名称
        config_map_name = choice_config['server_name']
        v1 = client.CoreV1Api()
        config_map = v1.read_namespaced_config_map(name=config_map_name, namespace=namespace)
        config_content = config_map.data.get(f'{config_map_name}.yaml')
        data_dict = yaml.safe_load(config_content)
        mongo_config = data_dict["mongo"]
        info = {
            "platform_type": choice_config["cloud_platform_environment"]["platform_type"],
            "mongo_uri": f"mongodb://{mongo_config['user']}:{mongo_config['password']}@{mongo_config['host']}/?authSource={mongo_config['authSource']}&authMechanism={mongo_config['authMechanism']}",
            "periodic_interval": data_dict["train"]["periodic_interval"],
            "initial_interval": data_dict["train"]["initial_interval"],
            "query_api": choice_config["cloud_platform_environment"]["query_api"],
            "upload_api": choice_config["cloud_platform_environment"]["upload_api"],
            "register_api": choice_config["cloud_platform_environment"]["register_api"],
            "deploy_api": choice_config["cloud_platform_environment"]["deploy_api"],
            "monitor_api": choice_config["cloud_platform_environment"]["monitor_api"],
            "database_name": choice_config["cloud_platform_environment"]['database_name'],
            "get_all_tenant_api": choice_config["cloud_platform_environment"]['get_all_tenant_api']
        }
    else:
        info = {
            "platform_type": choice_config["testing_environment"]["platform_type"],
            "mongo_uri": choice_config["testing_environment"]["mongo_uri"],
            "query_api": choice_config["testing_environment"]["query_api"],
            "upload_api": choice_config["testing_environment"]["upload_api"],
            "register_api": choice_config["testing_environment"]["register_api"],
            "deploy_api": choice_config["testing_environment"]["deploy_api"],
            "monitor_api": choice_config["testing_environment"]["monitor_api"],
            "periodic_interval": 259200,  # 3天
            "initial_interval": 1800,
            "database_name": "cement_ingredient",
            "get_all_tenant_api": choice_config["testing_environment"]['get_all_tenant_api']
        }
    return info


@app.on_event("startup")
async def startup_event():
    """
    服务初始化
    """
    global DATABASEINFO
    with open("config/config.yaml", "r", encoding="utf-8") as f:
        choice_config = yaml.safe_load(f)
    env_info = get_config_of_k8s(choice_config)
    DATABASEINFO = DatabaseInfo(
        platform_type=env_info["platform_type"],
        mongo_uri=env_info["mongo_uri"],
        periodic_interval=env_info["periodic_interval"],
        initial_interval=env_info["initial_interval"],
        query_api=env_info["query_api"],
        upload_api=env_info["upload_api"],
        register_api=env_info["register_api"],
        deploy_api=env_info["deploy_api"],
        monitor_api=env_info["monitor_api"],
        get_all_tenant_api=env_info["get_all_tenant_api"],
        headers=choice_config["headers"],
        parameters_for_3=choice_config["params_for_3d"],
        parameters_for_28=choice_config["params_for_28d"],
        mas_model_type=choice_config["mas_model_type"],
        train_optuna=choice_config["train_optuna"],
        model_type_suffix=choice_config["model_type_suffix"]
    )
    logging.info("相关配置信息:{}".format(DATABASEINFO.dict()))


@app.get("/v1/general_model_training")
@logging_decorator
async def train_main():
    """
    模型训练主入口
    """
    if DATABASEINFO.platform_type == 1:  # 表示在服务云上
        headers = DATABASEINFO.headers
        tenant_id_name_list = get_all_tenant_all(DATABASEINFO.get_all_tenant_api, headers)
    else:
        tenant_id_name_list = get_all_tenant_all(DATABASEINFO.get_all_tenant_api)
    logging.info("tenant_id_name_list: {}".format(tenant_id_name_list))
    # 获取全部的数据
    data_aggregator = DataProcessing(mongo_url=DATABASEINFO.mongo_uri, database_name=DATABASEINFO.database_name,
                                     tenant_id_list=tenant_id_name_list)
    all_data_df = data_aggregator.read_from_mongo_and_deal_data()

    if not os.path.exists(DATABASEINFO.save_path):
        os.makedirs(DATABASEINFO.save_path)

    training_main(data=all_data_df, save_model_path=DATABASEINFO.save_path, model_name=DATABASEINFO.model_name,model_desc_path = os.path.join(DATABASEINFO.save_path, "model_desc.json"),
               mas_model_type=DATABASEINFO.mas_model_type, params_for_3d=DATABASEINFO.parameters_for_3,
               params_for_28d=DATABASEINFO.parameters_for_28,
               train_optuna=DATABASEINFO.train_optuna,
               model_type_suffix=DATABASEINFO.model_type_suffix,
               mongo_url=DATABASEINFO.mongo_uri,
               database_name=DATABASEINFO.database_name,
               tenant_id_list=list(map(lambda x:x[0], tenant_id_name_list)))

    return {"tenant_id_name_list": tenant_id_name_list}

@app.post("/v1/download_model_file")
def get_download_model_file():
    """
    测试文件下载与续训练进行联合测试
    """
    file_path = os.path.join(DATABASEINFO.save_path, "CatBoost", "model.zip")
    return FileResponse(file_path, media_type="application/zip")


# 运行应用
if __name__ == "__main__":
    import uvicorn

    uvicorn.run("server_main:app", host="0.0.0.0", port=8080, reload=False)
